# -*- coding: utf-8 -*-
"""
PBAR_ExploreDetectorSets.py

Plot blocks of detectors.


Created on Mon May 05 08:28:14 2014

@author: jkwong
"""



#import csv
#import matplotlib.pyplot as plt
import os, cPickle, copy
import numpy as np
import matplotlib.pyplot as plt
#import numpy.matlib
#import datetime
#import time
import PBAR_Zspec
reload(PBAR_Zspec)
from sklearn.feature_selection import SelectKBest, chi2, f_classif, f_oneway
from sklearn.linear_model import LogisticRegression
from sklearn.lda import LDA
from sklearn import cross_validation
from sklearn.metrics import precision_recall_fscore_support
from sklearn.preprocessing import scale
from scipy.optimize import curve_fit

#from mpl_toolkits.mplot3d import Axes3D
# import mlpy
#from scipy.optimize import curve_fit

reload(PBAR_Zspec)

def gauss_function(x, a, x0, sigma):
    return a*np.exp(-(x-x0)**2/(2*sigma**2))

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2

######################################
##  LOAD DATA AND CALCULATE STATS   ##
######################################

# Threshold bin. The first "binThreshold" bins are set to zero
binThreshold = 8
# number of cross-validation cycles
numTrials = 50
# step size in count window
countStepSize = 25
# size of the window
windowSize = 50
# maximum count upper range window
MaxCount = 375

# the set number
setNum = 3

# SET THE DATASET LOCATION

# First set
if setNum == 1:
    basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data'
    infoFilename = os.path.join(basepath,'datasetSummaryOLD.txt')
# Second set
elif setNum == 2 or setNum == 3 or setNum == 4:
    # SECOND SET
    basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'
    infoFilename = os.path.join(basepath,'DatasetSummary2ndSet.txt')


# LOAD LIST OF FILES, DATASET GROUPS, LIST OF GOOD/BAD DETECTORS
if setNum == 1:
    # Create list of datasets
    (filenameList, fullfilenameList) = \
                   PBAR_Zspec.GenerateDefaultDatasetFilenameList(basepath)               
    # Create list of dataset groups
    (datasetGroups, datasetGroupsIndices) = \
                    PBAR_Zspec.GenerateDefaultDatasetGroupList(filenameList)
    # Create list of good/bad detectors
    (goodDetectorsList, badDetectorsList, goodDetectorsMask, badDetectorsMask) = \
                        PBAR_Zspec.GenerateDefaultDetectorList()
    (goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
            goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists()
elif setNum == 2 or setNum == 3 or setNum == 4:
    # Create list of datasets
    (filenameList, fullfilenameList) = \
                   PBAR_Zspec.GenerateDefaultDatasetFilenameList2ndSet(basepath)               
    # Create list of dataset groups
    (datasetGroups, datasetGroupsIndices, datasetGroupsWidth, datasetRawWidth, datasetMaterial) = \
                    PBAR_Zspec.GenerateDefaultDatasetGroupList2ndSet(filenameList)
    # Create list of good/bad detectors
    (goodDetectorsList, badDetectorsList, goodDetectorsMask, badDetectorsMask) = \
                        PBAR_Zspec.GenerateDefaultDetectorList2ndSet()
    (goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
            goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists2ndSet()

# LOAD DATASET SUMMARY
(datasetDescription, datasetAcquisitionTime, \
    datasetTime, datasetTimeNum, datasetTimeStr) = \
    PBAR_Zspec.GetDatasetInformation(infoFilename, filenameList)

# LOAD ZSPEC DATA
print "Loading Data"
dat = PBAR_Zspec.ReadZspec(fullfilenameList)

# Load Radiography data
#(datRad, datRadZspec, radMap) = PBAR_Zspec.ReadRad(basepath)

# Read in the Calibration file
(calTimeNum, calGainShift) = PBAR_Zspec.LoadGainCalibration(os.path.join(basepath, 'GainCorrectionVer2.csv'))

# Generate extrapolated gain matrix
# misnomer - is actually the gain calibration bin
gainExtrapolated = PBAR_Zspec.ExtrapolateGain(calTimeNum, calGainShift, datasetTimeNum)

# Multiple the bin array with this to calibrate
gainCorrection = 100.0/gainExtrapolated

# Xray Pulse Rate
if setNum == 1:
    # Pulse rate which has always be 60 Hz
    pulseRate = 60 * np.ones(len(dat))
elif setNum == 2 or setNum == 3 or setNum == 4:
    # Pulse rate is not always 60Hz unlike in the first set.
    filenameList = list(filenameList)
    pulseRate = np.ones(len(filenameList)) * 250.
    pulseRate[0] = 60.
    pulseRate[filenameList.index('ec76')] = 60.
    pulseRate[filenameList.index('ec84')] = 60.
    pulseRate[filenameList.index('ec85')] = 60.
    filenameList = np.array(filenameList)

# calculate the stats
#stats = PBAR_Zspec.CalculateStats(dat, datasetAcquisitionTime, np.arange(256).astype(float), gainCorrection, basepath, binThreshold, pulseRate)
#stats = PBAR_Zspec.CalculateStats(dat, datasetAcquisitionTime, np.arange(256).astype(float), gainCorrection, binThreshold, pulseRate)
# Interpolation version to lessen discrete binning effects
stats = PBAR_Zspec.CalculateStatsInterp(dat, datasetAcquisitionTime, np.arange(256).astype(float), gainCorrection, binThreshold, pulseRate)

## Save values to file
#PBAR_Zspec.WriteStats(stats, filenameList, outputDir, 'd')

######################################################
##  MAKE ARRAYS FOR FEATURES SELECTION FUNCTION


# Detectors worth including
detectorList0 = goodDetectorsList[(goodDetectorsList >= 41) & (goodDetectorsList < 120)]
detectorList1 = goodDetectorsList[(goodDetectorsList >= 31) & (goodDetectorsList < 60)]
detectorList2 = goodDetectorsList[(goodDetectorsList >= 61) & (goodDetectorsList < 90)]
detectorList3 = goodDetectorsList[(goodDetectorsList >= 91) & (goodDetectorsList < 120)]


# The dict containing the collapsed stuff will contain only these groups of data sets
#groupNamesList = np.array(('PbALL', 'Pb', 'Fe', 'Al', 'PbNOT', \
#                           'PbALLALL', 'PbNOTNOT', \
#                           'Pb3Fe', 'Pb3Al', 'Pb4Fe', 'Pb4Al', 'Pb5Fe', 'Pb5Al'))
#                           
#groupNamesList = np.array(('PbALL', 'Pb', 'Fe', 'Al', 'PbNOT', \
#                           'PbALLALL', 'PbNOTNOT', 'PbLS', 'FeLS', 'AlLS', \
#                           'Pb3Fe', 'Pb3Al', 'Pb4Fe', 'Pb4Al'))

if setNum == 1:  # set1 doesn't have Pb5Al or Pb5Fe
    groupNamesList = np.array(('PbALL', 'Pb', 'Fe', 'Al', 'PbNOT', \
                               'PbALLALL', 'PbNOTNOT', \
                               'Pb3Fe', 'Pb3Al', 'Pb4Fe', 'Pb4Al'))
elif setNum == 2 or setNum == 3 or setNum == 4:
    groupNamesList = np.array(('PbALL', 'Pb', 'Fe', 'Al', 'PbNOT', \
                           'PbALLALL', 'PbNOTNOT', \
                           'Pb3Fe', 'Pb3Al', 'Pb4Fe', 'Pb4Al', 'Pb5Fe', 'Pb5Al', \
                           'PbALL_0', 'PbALL_1', 'PbNOT_0', 'PbNOT_1'))

# Make a dictionary of dictionaries containing a values from a subset of the detectors.
# statsCollapsed[<dataset group name>][<statistics name>] = 
#        values for all detectors specified for all datasets in the group
#
statsCollapsed0 = PBAR_Zspec.CreateStatsCollapsed(stats, groupNamesList, datasetGroupsIndices, detectorList0)
statsCollapsed1 = PBAR_Zspec.CreateStatsCollapsed(stats, groupNamesList, datasetGroupsIndices, detectorList1)
statsCollapsed2 = PBAR_Zspec.CreateStatsCollapsed(stats, groupNamesList, datasetGroupsIndices, detectorList2)
statsCollapsed3 = PBAR_Zspec.CreateStatsCollapsed(stats, groupNamesList, datasetGroupsIndices, detectorList3)


# This defines the columns of the feature matrix (X)
#statsCollapseListALL = np.array(( \
#    'binSTD_binMean', \
#    'binMean_g1', 'binSTD_g1', \
##    'binMean', 'binSTD', \
#    'binMean_q50_g1', 'q_range_ratio',  \
#    'multibin_20_ratio_g1', 'multibin_10_ratio_g1'))

statsCollapseListALL = np.array(( \
    'binSTD_binMean', \
#    'binMean', 'binSTD', \
    'binMean_g1', 'binSTD_g1', \
    'binMean_q50_g1', 'q_range_ratio',  \
    'multibin_20_ratio_g1', 'multibin_10_ratio_g1', \
    'q_range_g1', 'q_50_g1',
    'binMoment3', 'binMoment4', 'binSkew', 'binKur'
    ))

statsCollapseListALL = np.array(( \
    'binSkew', 'binKur', \
    'binMean_g1', 'binSTD_g1', \
    'binMean', \
#    'q_range_ratio',  \
    'multibin_20_ratio_g1', \
#    'q_range_g1', 'q_50_g1',
    ))


# Make the stats matrix that contains all the features we want
# Only include data from these groups for the test/training matrix
# groupNamesExportList = np.array(('PbALL', 'PbNOT'))

if setNum == 1 or setNum == 2:
    groupNamesExportList = np.array(('PbNOT', 'PbALL'))
elif setNum == 3:
    groupNamesExportList = np.array(('PbNOT_0', 'PbALL')) # assumes that the Pb is fine in 2a and 2b

elif setNum == 4:
    groupNamesExportList = np.array(('PbNOT_1', 'PbALL')) # assumes that the Pb is fine in 2a and 2b

# - statsMatrix contains the features and truth value
# - statsCollapsedMatrix, dict of two matrices, one containing only Pb, other containing not Pb
# statsMatrix, vertical concatenation of the two matrices in statsCollapsedMatrix; the last colum is the target value (1 = lead)
# make matrix containing only the counts
# - countCollapsedMatrix - same as statsCollapsedMatrix except only contain counts
# - countMatrix - same as statsMatrix but only contains counts
(statsCollapsedMatrix0, statsMatrix0) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed0, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix0, countMatrix0) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed0, ['count'], groupNamesExportList)
countMatrix0 = countMatrix0[:,0]

(statsCollapsedMatrix1, statsMatrix1) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed1, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix1, countMatrix1) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed1, ['count'], groupNamesExportList)
countMatrix1 = countMatrix1[:,0]

(statsCollapsedMatrix2, statsMatrix2) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed2, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix2, countMatrix2) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed2, ['count'], groupNamesExportList)
countMatrix2 = countMatrix2[:,0]

(statsCollapsedMatrix3, statsMatrix3) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed3, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix3, countMatrix3) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed3, ['count'], groupNamesExportList)
countMatrix3 = countMatrix3[:,0]


# Plot


# plot of features vs counts
# pb, not pb
# 

featureIndex = np.where(statsCollapseListALL == 'binSkew')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binMean')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binSkew')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binKur')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binMean_g1')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binKur')[0][0]
#featureIndex = 12

plotPbNOT = 0
plotPbAll = 1

plt.figure()
plt.grid()

## all lead, all no lead
#cut = statsMatrix0[:,-1] == 0
#plt.plot(countMatrix0[cut], statsMatrix0[cut,featureIndex], 'sb', markersize = 7, alpha = 0.2, label = 'Not Pb, 41-120')
#cut = statsMatrix0[:,-1] == 1
#plt.plot(countMatrix0[cut], statsMatrix0[cut,featureIndex], 'sr', markersize = 7, alpha = 0.2, label = 'Pb, Set 41-120')
##
if plotPbNOT:
    cut = statsMatrix1[:,-1] == 0
    plt.plot(countMatrix1[cut], statsMatrix1[cut,featureIndex], '^b', markersize = 7, alpha = 0.3, label = 'Not Pb, 31-60')
if plotPbAll:
    cut = statsMatrix1[:,-1] == 1
    plt.plot(countMatrix1[cut], statsMatrix1[cut,featureIndex], '^r', markersize = 7, alpha = 0.3, label = 'Pb, Set 31-60')

if plotPbNOT:
    cut = statsMatrix2[:,-1] == 0
    plt.plot(countMatrix2[cut], statsMatrix2[cut,featureIndex], 'dc', markersize = 7, alpha = 0.3, label = 'Not Pb, 61-90')
if plotPbAll:
    cut = statsMatrix2[:,-1] == 1
    plt.plot(countMatrix2[cut], statsMatrix2[cut,featureIndex], 'dm', markersize = 7, alpha = 0.3, label = 'Pb, Set 61-90')

if plotPbNOT:
    cut = statsMatrix3[:,-1] == 0
    plt.plot(countMatrix3[cut], statsMatrix3[cut,featureIndex], 'sy', markersize = 7, alpha = 0.3, label = 'Not Pb, 91-120')
if plotPbAll:
    cut = statsMatrix3[:,-1] == 1
    plt.plot(countMatrix3[cut], statsMatrix3[cut,featureIndex], 'sg', markersize = 7, alpha = 0.3, label = 'Pb, Set 91-120')


plt.xlabel('Count Per Pulse')
plt.ylabel(statsCollapseListALL[featureIndex])
#plt.legend(loc =1)
plt.legend(prop={'size':16})

#if featureIndex == 9:
#    axis((0, 400, 0, 3e4))
#elif featureIndex == 11:
#    axis((0, 400, 0, 2))
#elif featureIndex == 12:
#    axis((0, 400, 2, 7))
    
#axis((0, 400, 3, 7))
#axis((0, 400, 10, 50))

